import os
import pandas as pd
import time
from datetime import datetime, timedelta  
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures import ThreadPoolExecutor, as_completed
import traceback
from tqdm import tqdm


start_date = datetime(2024, 1, 15)   
end_date = datetime.today()  

def mysql_engine(host, port, user, passwd, db=None):
    from sqlalchemy import create_engine
    try:
        engine = create_engine(f'mysql+pymysql://{user}:{passwd}@{host}:{port}/{db}')
        return engine
    except Exception as e:
        print(f"An error occurred: {e}")
mysql_rule = mysql_engine('47.253.56.86',4001,'longxiaolei','WWsCQ9TXG+BJBHlK','rule')


def task(sql,date_string,pbar):
    try:
        start_time1 = time.time()
        df = pd.read_sql(sql,mysql_rule)
        end_time1 = time.time()
        print(f"{date_string},{df.shape},查询耗时{end_time1-start_time1}")
        merchant_ids = df['merchant_id'].unique()
        types  = df['type'].unique()
        len_arr = [0]
        for merchant_id in merchant_ids : 
            for t in types:
                # 注意在savepath中设定保存的路径和文件名
                df_new = df[(df['merchant_id']==merchant_id)&(df['type']==t)]
                if df_new.shape[0]>0:
                    len_arr.append(df_new.shape[0])
                    save_path = f"data/{merchant_id}_type{t}_{date_string}.parquet"
                    df_new.to_parquet(save_path,compression='zstd')
                    print(save_path,df_new.shape)
        if sum(len_arr ) == df.shape[0]:
            print('通过基本数据检查========>')
        else:
            print("^^^^^^^^^^^^^数据条数检查失败^^^^^^^^^^^^^")
        time.sleep(0.5)
        pbar.update(1)
    except Exception as e:
        print(f"多线程任务出现了异常，信息如下{e}")
        traceback.print_exc()

        
if __name__ == '__main__':
    start_time = time.time()    
    single_days = []
    for single_day in (start_date + timedelta(days=x) for x in range(0, (end_date - start_date).days + 1)):  
        single_days.append(single_day)
        
    pool = ThreadPoolExecutor(15)
    with (tqdm(total=len(single_days), desc="Downloading") as pbar):
        all_task_result = []
        for single_day in single_days:
            # 格式化日期为字符串  
            date_string = single_day.strftime("%Y-%m-%d")  
            next_day = single_day+timedelta(days=1) 
            next_day_str = next_day.strftime("%Y-%m-%d")  
            
            sql = f"""select a.id,a.merchant_id,a.tx_id,a.req_data,a.resp_data,a.create_time,a.update_time, 
                    case when a.type is null then b.type else a.type end as type
            from risk.t_risk_urule_req_log a left join risk.t_risk_req_record b on a.tx_id = b.tx_id and a.merchant_id = b.merchant_id
            where a.create_time >= '{date_string}' and a.create_time < '{next_day_str}' and a.tx_id is not null
            """
            
            all_task_result.append(pool.submit(task,sql,date_string,pbar))
        for task_result in as_completed(all_task_result):
            msg = task_result.result()

    end_time = time.time()
    print(f"本次计算结束，共计算 {len(single_days)}的数据，耗时 {round(end_time - start_time, 3)} s")

        










    

    